// Tencent is pleased to support the open source community by making ncnn available.
//
// Copyright (C) 2021 THL A29 Limited, a Tencent company. All rights reserved.
//
// Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
// in compliance with the License. You may obtain a copy of the License at
//
// https://opensource.org/licenses/BSD-3-Clause
//
// Unless required by applicable law or agreed to in writing, software distributed
// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
// CONDITIONS OF ANY KIND, either express or implied. See the License for the
// specific language governing permissions and limitations under the License.

#include "seg.h"

#include <opencv2/core/core.hpp>
#include <opencv2/imgproc/imgproc.hpp>

#include "cpu.h"

#define MAX_STRIDE 32



SegNail::SegNail()
{
    blob_pool_allocator.set_size_compare_ratio(0.f);
    workspace_pool_allocator.set_size_compare_ratio(0.f);
}


int SegNail::load(AAssetManager* mgr, bool use_gpu)
{
    seg.clear();
    blob_pool_allocator.clear();
    workspace_pool_allocator.clear();

    ncnn::set_cpu_powersave(2);
    ncnn::set_omp_num_threads(ncnn::get_big_cpu_count());

    seg.opt = ncnn::Option();

#if NCNN_VULKAN
    seg.opt.use_vulkan_compute = use_gpu;
#endif

    seg.opt.num_threads = ncnn::get_big_cpu_count();
    seg.opt.blob_allocator = &blob_pool_allocator;
    seg.opt.workspace_allocator = &workspace_pool_allocator;

//    char parampath[256];
//    char modelpath[256];
//    sprintf(parampath, "yolov8%s-seg-sim-opt-fp16.param", modeltype);
//    sprintf(modelpath, "yolov8%s-seg-sim-opt-fp16.bin", modeltype);

    seg.load_param(mgr,"segnail-sim-opt.param");
    seg.load_model(mgr,"segnail-sim-opt.bin");

    mean_vals[0] = 0.5*255.0;
    mean_vals[1] = 0.5*255.0;
    mean_vals[2] = 0.5*255.0;
    norm_vals[0] = 1 / 127.5;
    norm_vals[1] = 1 / 127.5;
    norm_vals[2] = 1 / 127.5;

    return 0;
}

int SegNail::detect(const cv::Mat& rgb)
{
    int width = rgb.cols;
    int height = rgb.rows;


    // 0是背景，1是前景；背景不上色，前景蓝色
    // 分割要注意输入尺寸；512可能会出现黑条，416实测可行；F.interpolate corner不要对齐

    ncnn::Mat in = ncnn::Mat::from_pixels_resize(rgb.data, ncnn::Mat::PIXEL_BGR, width, height, input_w, input_h);

    // pad to target_size rectangle

    in.substract_mean_normalize(mean_vals, norm_vals);

    ncnn::Extractor ex = seg.create_extractor();

    ex.input("input", in);

    ncnn::Mat out;
    ex.extract("output", out);

    LOGD("hhhhh h:%d,w:%d",height, width);

    cv::Mat mask(out.h, out.w, CV_8UC1);
    const float* maskmap0 = out.channel(0);
    const float* maskmap1 = out.channel(1);

    for (int i{ 0 }; i < out.h; i++) {
        for (int j{ 0 }; j < out.w; ++j) {
            mask.at<uchar>(i, j) = maskmap1[i * out.w + j] > maskmap0[i * out.w + j] ? 255 : 0;
        }
    }
    ncnn::Mat ncnn_mask = ncnn::Mat::from_pixels_resize(mask.data, ncnn::Mat::PIXEL_GRAY, mask.cols, mask.rows,rgb.cols, rgb.rows);

    cv::Mat result(ncnn_mask.h, ncnn_mask.w, CV_8UC1);
    ncnn_mask.to_pixels(result.data, ncnn::Mat::PIXEL_GRAY);


    cv::Mat yellowImg(result.size(), CV_8UC3, cv::Scalar(0, 0, 255)); // 蓝色
    // 复制原始图像到结果图像中，只在掩码为 1 的地方添加透明的黄色效果
    cv::Mat clone_img = rgb.clone(); // 克隆原始图像
    cv::Mat mask_3c; // 将单通道掩码转换为三通道

    cv::cvtColor(result, mask_3c, cv::COLOR_GRAY2BGR);
    cv::Mat masked_yellow = yellowImg.mul(mask_3c); // 在掩码为 1 的地方应用透明的黄色效果
    cv::addWeighted(clone_img, 1.0, masked_yellow, 1.0, 0.0, clone_img);
    clone_img.copyTo(rgb);


    return 0;
}
